##
## Copyright (C) 2009-2012 Luciano Xumerle <luciano.xumerle@gmail.com>
##
## This program is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.
##
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with this program.  If not, see <http://www.gnu.org/licenses/>.
##

import re, string
import time

import sys
reload(sys)
sys.setdefaultencoding("latin1")


def plog( msg ):
   print >>sys.stderr, time.ctime() + ' | ' + msg

## ELAND ++ ELAND ++ ELAND ++ ELAND ##
## ELAND ++ ELAND ++ ELAND ++ ELAND ##
## ELAND ++ ELAND ++ ELAND ++ ELAND ##

def splitEland2line(line, onlyMulti):
   """Returns a list of list with elements:
      sequence|chromosome|direction|sstart|sstop|mismatches|full_uc_code
      onlyMulti==False -> use unique match only
      onlyMulti==True -> use multi matches only"""
   process=re.split('\s', line.strip())
   if len(process)<4:
       return []
   mmm=re.split(':', process[2])
   isMultiMatch=(int(mmm[0])+int(mmm[1])+int(mmm[2]))>1
   res=[]
   if onlyMulti == isMultiMatch:
       oldchr=''
       matches=re.split(',', process[3])
       for mm in matches:
	   twodot=mm.find(':')
           if twodot>0:
               chr=mm[0:twodot]
               chr=chr[0:chr.find(".fa")]
               if chr.find("splice")>-1:
                       chr=mm[mm.find('/')+1:mm.find('|')]
               oldchr=chr
           else:
               chr=oldchr
               twodot=-1
	   verse=mm.rfind('F')
	   if verse<0:
	       verse=mm.rfind('R')
	   mismatch=mm[verse+1:]
	   sstart=mm[twodot+1:verse]
	   verse=mm[verse]
	   sstop=int(sstart)+len(process[1])
	   row=[]
           row.append(process[0])
	   row.append(chr)
	   row.append(verse)
	   row.append(sstart)
	   row.append(sstop)
	   row.append(mismatch)
           row.append(process[2][mm.find('/')+1:]) # the full uc code if needed
           res.append(row)
       return res
   return []


## BOWTIE ++ BOWTIE ++ BOWTIE ++ BOWTIE ++ BOWTIE ++ BOWTIE ##
## BOWTIE ++ BOWTIE ++ BOWTIE ++ BOWTIE ++ BOWTIE ++ BOWTIE ##
## BOWTIE ++ BOWTIE ++ BOWTIE ++ BOWTIE ++ BOWTIE ++ BOWTIE ##
## NB: bowtie starts to count bases from zero, eland from 1 ##
## NB: bowtie starts to count bases from zero, eland from 1 ##

def splitBowtieLine(line):
   """Returns a list with elements:"""
   """sequence|chromosome|direction|sstart|sstop|mismatches|full_uc_code"""
   process=re.split('\s+', line.strip())
   if len(process)<7:
       return []
   mismatches=0
   if len(process)==8:
       mismatches=len( process[7].split(',') )
   row=[]
   row.append(process[0]) # name
   if process[2].find('uc')>-1:
       row.append( process[2][:process[2].find('|')] )
   else:
       row.append(process[2]) # chr
   row.append(process[1]) # verse
   row.append(int(process[3])+1) # start
   row.append(int(process[3])+len(process[4])) # stop
   row.append(mismatches) # mismatches
   row.append(process[2]) # the full uc code if needed
   return row


## maq align == eland align
def splitMaqLine(line):
   """Returns a list with elements:"""
   """sequence|chromosome|direction|sstart|sstop|mismatches|full_uc_code"""
   process=re.split('\t+', line.strip())
   if len(process)<14:
       return []
   row=[]
   row.append(process[0]) # name
   if process[1].find('|')>-1:
       row.append( process[1][:process[1].find('|')] )
   else:
       row.append(process[1]) # chr
   row.append(process[3]) # verse
   row.append(int(process[2])) # start
   row.append(int(process[2])+int(process[13])) # stop
   row.append(process[9]) # mismatches
   row.append(process[1]) # the full uc code if needed
   return row


## split samtools pileup > samtools  pileup -f genome.fa assay.sorted.bam
def splitPileupLine( line ):
   """Returns a list with elements:
      chrom|position|coverage|sequence|quality"""
   rr=line.strip().split("\t")
   return [ rr[0], int(rr[1]), rr[2], int(rr[3]), rr[4].upper(), rr[5] ]


##
## PARSE PILEUP ROW
## line: (0) chr1 (1) 14542 (2) A (3) 24 (4) C..G (5) ##F?
##
def parsePileupRow( line, assay ) :
   """Parse pileup row and returns alleles with the reads count.
      Output list is [ [ assay, chrom, pos, ref, count, var1, count1, var2...  ] ]
      Variation may be not present. Last element in list is reads total count."""
   row=re.split("\t", line.strip())
   chrom=row[0]
   pos=row[1]
   ref=row[2].upper()
   count=int(row[3])
   string=row[4]
   alleles=dict()
   alleles[ref]=0

   string = string.upper()
   string = pileupInsDel( string, "+" )
   string = pileupInsDel( string, "-" )

   for a in string:
       if a == "." or a == "," :
           alleles[ref] = alleles[ref] + 1
       elif a in "ACTG+-" :
           if not a in alleles.keys():
               alleles[a] = 0
           alleles[a] = alleles[a] + 1

   res=[ assay, chrom, pos, ref, alleles[ref] ]
   if len(alleles.keys())<2 :
       return res

   table=[]
   for a in alleles.keys() :
       if a == ref:
           continue
       table.append( [ a, alleles[a] ] )
   table.sort(key=lambda x: x[1], reverse=True )

   res=[ assay, chrom, pos, ref, alleles[ref] ]
   for a in table :
       res.append( a[0] )
       res.append( a[1] )
   res.append( count )
   return res

##
## MANAGE PILEUP INSDEL
## (used by parsePileupRow)
##
def pileupInsDel ( pil, insdel ):
   """Search for insdel in pileup string. Parameter insdel id + or -"""
   if not insdel in pil:
       return pil
   res = ""
   snum = ""
   num=0
   string = ""
   do = False
   if not insdel in pil:
       return pil
   for a in pil:
       if a == insdel:
           res = res + insdel
           do = True
       elif do:
           if a.isdigit():
               snum = snum + a
           else:
               num = int(num)-1
               snum=""
               do = False
       elif num>0:
           num=num-1
       else:
           res = res + a
   return res


## split samtools SAM > samtool view assay.sorted.bam chr1:100-200
def splitSamLine( line ):
   """Returns a list with elements:
      chrom|position|length|sequence|quality|id"""
   rr=line.strip().split("\t")
   return [ rr[2], int(rr[3]), len(rr[9]), rr[9].upper(), rr[10], rr[0] ]


## read newregions.good file created from erange
def readNewRegionsGood( nrgfile ):
   """Returns a list with row. Each row contains fields:
      chrom|start|stop"""
   res = []
   for line in open(nrgfile, "r"):
       row=line.strip().split("\t")
       if "chr" != row[1][:3]:
           continue
       chrom=row[1]
       start=int(row[2])
       stop=int(row[3])
       res.append( [chrom, start, stop] )
   return res

